#include "keygen.h"

KeyGenerator::KeyGenerator(Context *parameters) : parameters(parameters)
{
    KeyedPRNG *prng = new KeyedPRNG();
    poolQ = parameters->ringQ->NewPoly();
    poolQP = parameters->NewRingQP()->NewPoly();
    gaussianSamplerQ = new GaussianSampler(prng, parameters->ringQ, parameters->sigma, int(6 * parameters->sigma));
    uniformSamplerQ = new UniformSampler(prng, parameters->ringQ);
    uniformSamplerP = new UniformSampler(prng, parameters->ringP);
}

std::pair<SecretKey *, PublicKey *> KeyGenerator::GenKeyPair()
{
    SecretKey *sk = GenSecretKey();
    return { sk, GenPublicKey(sk) };
}

SecretKey *KeyGenerator::GenSecretKey()
{
    return GenSecretkeyWithDistrib(1.0 / 3);
}
/*
 * @brief 生成秘密密钥
 *
 * @param p 概率
 * @return SecretKey* 返回生成的秘密密钥
 *
 * 该函数:
 * - 创建一个新的KeyedPRNG
 * - 创建一个新的TernarySampler
 * - 调用genSecretKeyFromSampler生成秘密密钥
 */
SecretKey *KeyGenerator::GenSecretkeyWithDistrib(double p)
{
    KeyedPRNG *prng = new KeyedPRNG();
    TernarySampler *ternarySampler = new TernarySampler(prng, parameters->ringQ, p);
    return genSecretKeyFromSampler(ternarySampler);
}
/*
 * @brief 生成公钥
 *
 * @param sk 输入密钥
 * @return PublicKey* 返回生成的公钥
 *
 * 该函数:
 * - 创建一个新的公钥
 * - 读取采样器的数据
 * - 将数据扩展到QP多项式中
 * - 进行NTT变换
 * - 返回生成的公钥
 */
PublicKey *KeyGenerator::GenPublicKey(SecretKey *sk)
{
    PublicKey *pk = new PublicKey(parameters);

    int levelQ = parameters->QCount() - 1;
    int levelP = parameters->PCount() - 1;

    // pk[0] = [-as + e]
    // pk[1] = [a]
    gaussianSamplerQ->Read(pk->value[0].Q);
    parameters->NewRingQP()->ExtendBasisSmallNormAndCenter(pk->value[0].Q, levelP, nullptr, pk->value[0].P);

    parameters->ringQ->NTTLvl(levelQ, pk->value[0].Q, pk->value[0].Q);
    parameters->ringP->NTTLvl(levelP, pk->value[0].P, pk->value[0].P);

    uniformSamplerQ->Read(pk->value[1].Q);
    uniformSamplerP->Read(pk->value[1].P);

    parameters->ringQ->MulCoeffsBarrettAndSubLvl(levelQ, sk->value.Q, pk->value[1].Q, pk->value[0].Q);
    parameters->ringP->MulCoeffsBarrettAndSubLvl(levelP, sk->value.P, pk->value[1].P, pk->value[0].P);

    return pk;
}
/*
 * @brief 从采样器生成秘密密钥
 *
 * @param sampler 采样器
 * @return SecretKey* 返回生成的秘密密钥
 *
 * 该函数:
 * - 创建一个新的秘密密钥
 * - 读取采样器的数据
 * - 将数据扩展到QP多项式中
 * - 进行NTT变换
 * - 返回生成的秘密密钥
 */
SecretKey *KeyGenerator::genSecretKeyFromSampler(BaseSampler *sampler)
{
    SecretKey *sk = new SecretKey(parameters);
    int levelQ = parameters->QCount() - 1;
    int levelP = parameters->PCount() - 1;
    sampler->Read(sk->value.Q);
    parameters->NewRingQP()->ExtendBasisSmallNormAndCenter(sk->value.Q, levelP, nullptr, sk->value.P);

    parameters->ringQ->NTTLvl(levelQ, sk->value.Q, sk->value.Q);
    parameters->ringP->NTTLvl(levelP, sk->value.P, sk->value.P);

    return sk;
}
/*
 * @brief 生成旋转密钥集
 *
 * @param ks 旋转次数列表
 * @param includeConjugate 是否包含共轭旋转
 * @param sk 输入密钥
 * @return RotationKeySet* 返回创建的旋转密钥集
 *
 * 该函数:
 * - 创建一个新的旋转密钥集
 * - 遍历给定的旋转次数列表
 * - 为每个旋转次数生成伽罗瓦元素
 * - 将这些伽罗瓦元素存储在galEls容器中
 * - 如果includeConjugate为true,则添加共轭旋转的伽罗瓦元素
 * - 调用GenRotationKeys生成旋转密钥集
 */
RotationKeySet *KeyGenerator::GenRotationKeysForRotations(std::vector<int> ks, bool includeConjugate, SecretKey *sk)
{
    std::vector<uint64_t> galEls(ks.size());
    for (int i = 0; i < ks.size(); i++) {
        galEls[i] = parameters->GaloisElementForColumnRotationBy(ks[i]);
    }
    if (includeConjugate) {
        galEls.push_back(parameters->GaloisElementForRowRotation());
    }
    return GenRotationKeys(galEls, sk);
}
/*
 * @brief 根据galois元素创建并返回密钥
 *
 * @param galEls 伽罗瓦元素列表
 * @param sk 输入密钥
 * @return GaloisKeys* 返回创建的伽罗瓦密钥
 *
 * 该函数:
 * - 创建一个新的伽罗瓦密钥
 * - 遍历给定的伽罗瓦元素列表
 * - 为每个伽罗瓦元素生成旋转密钥
 * - 将这些旋转密钥存储在rks->keys容器中
 */
GaloisKeys *KeyGenerator::GenGaloisKeys(std::vector<uint64_t> &galEls, SecretKey *sk)
{
    GaloisKeys *rks = new GaloisKeys(parameters, galEls);
    for (int i = 0; i < galEls.size(); i++) {
        uint64_t galEl = galEls[i];
        GenrotKey(sk->value, parameters->InverseGaloisElement(galEl), rks->keys[galEl]);
    }
    return rks;
}

/*
 * @brief 生成旋转密钥集
 *
 * @param galEls 伽罗瓦元素列表
 * @param sk 输入密钥
 * @return RotationKeySet* 返回创建的旋转密钥集
 *
 * 该函数:
 * - 创建一个新的旋转密钥集
 * - 遍历给定的伽罗瓦元素列表
 * - 为每个伽罗瓦元素生成旋转密钥
 * - 将这些旋转密钥存储在rks->keys容器中
 */
RotationKeySet *KeyGenerator::GenRotationKeys(std::vector<uint64_t> &galEls, SecretKey *sk)
{
    RotationKeySet *rks = new RotationKeySet(parameters, galEls);
    for (int i = 0; i < galEls.size(); i++) {
        uint64_t galEl = galEls[i];
        GenrotKey(sk->value, parameters->InverseGaloisElement(galEl), rks->keys[galEl]);
    }
    return rks;
}

/*
 * @brief 生成旋转切换密钥
 *
 * @param sk 输入密钥
 * @param galEl 伽罗瓦元素
 * @param swk 切换密钥
 *
 * 该函数:
 * - 计算Q和P基的层级
 * - 计算P基的模数
 * - 遍历每个分解块,生成切换密钥
 */
void KeyGenerator::GenrotKey(PolyQP sk, uint64_t galEl, SwitchingKey *swk)
{
    std::vector<uint64_t> index = parameters->ringQ->PermuteNTTIndex(galEl);
    parameters->ringQ->PermuteNTTWithIndexLvl(parameters->QCount() - 1, sk.Q, index, poolQP->Q);
    parameters->ringQ->PermuteNTTWithIndexLvl(parameters->PCount() - 1, sk.P, index, poolQP->P);

    PrivateGenSwitchingKey(sk.Q, poolQP, swk);
}
/*
 * @brief 生成切换密钥
 *
 * @param skIn 输入密钥
 * @param skOut 输出密钥
 * @param swk 切换密钥
 *
 * 该函数:
 * - 计算Q和P基的层级
 * - 计算P基的模数
 * - 遍历每个分解块,生成切换密钥
 */
void KeyGenerator::PrivateGenSwitchingKey(Poly *skIn, PolyQP *skOut, SwitchingKey *swk)
{
    int levelQ = swk->value[0][0]->Q->Coeffs.size() - 1;
    int levelP = swk->value[0][0]->P->Coeffs.size() - 1;

    mpz_class pBigInt;
    if (levelP == parameters->PCount() - 1) {
        pBigInt = *(parameters->ringP->modulusBigint);
    } else {
        std::vector<uint64_t> &P = parameters->ringP->Modulus;
        pBigInt = P[0];
        for (auto i = 1; i < levelP + 1; i++) {
            pBigInt *= P[i];
        }
    }
    // Computes P * skIn
    parameters->ringQ->MulScalarBigintLvl(levelQ, skIn, &pBigInt, poolQ);

    int alpha = levelP + 1;
    int beta = int(ceil(double(levelQ + 1) / double(levelP + 1)));

    int index;
#pragma omp parallel for
    for (int i = 0; i < beta; i++) {
        // e
        gaussianSamplerQ->ReadLvl(levelQ, swk->value[i][0]->Q);
        parameters->NewRingQP()->ExtendBasisSmallNormAndCenter(swk->value[i][0]->Q, levelP, nullptr,
            swk->value[i][0]->P);
        parameters->NewRingQP()->NTTLazyLvl(levelQ, levelP, swk->value[i][0], swk->value[i][0]);

        // a (since a is uniform, we consider we already sample it in the NTT and Montgomery domain)
        uniformSamplerQ->ReadLvl(levelQ, swk->value[i][1]->Q);
        uniformSamplerP->ReadLvl(levelP, swk->value[i][1]->P);

        // e + (skIn * P) * (q_star * q_tild) mod QP
        //
        // q_prod = prod(q[i*alpha+j])
        // q_star = Q/qprod
        // q_tild = q_star^-1 mod q_prod
        //
        // Therefore : (skIn * P) * (q_star * q_tild) = sk*P mod q[i*alpha+j], else 0
        for (int j = 0; j < alpha; j++) {
            index = i * alpha + j;
            // It handles the case where nb pj does not divide nb qi
            if (index >= levelQ + 1) {
                break;
            }
            uint64_t qi = parameters->ringQ->Modulus[index];
            std::vector<uint64_t> &p0tmp = poolQ->Coeffs[index];
            std::vector<uint64_t> &p1tmp = swk->value[i][0]->Q->Coeffs[index];
            for (uint64_t w = 0; w < parameters->ringQ->N; w++) {
                p1tmp[w] = CRed(p1tmp[w] + p0tmp[w], qi);
            }
        }
        // (skIn * P) * (q_star * q_tild) - a * skOut + e mod QP
        parameters->NewRingQP()->MulCoeffsBarrettAndSubLvl(levelQ, levelP, swk->value[i][1], skOut, swk->value[i][0]);
    }
}

/*
 * @brief 创建重线性化密钥
 *
 * @param sk 输入密钥
 * @param maxDegree 最大重线性化度
 * @return RelinearizationKey* 返回创建的重线性化密钥
 *
 * 该函数:
 * - 检查P基是否为空
 * - 创建一个新的重线性化密钥
 * - 将输入密钥的Q基系数复制到QP多项式中
 * - 遍历最大重线性化度,生成重线性化密钥
 */
RelinearizationKey *KeyGenerator::GenRelinearizationKey(SecretKey *sk, int maxDegree)
{
    if (parameters->PCount() == 0) {
        throw std::runtime_error("modulus P is empty");
    }

    RelinearizationKey *rlk = new RelinearizationKey(parameters, maxDegree);
    poolQP->Q->CopyValues(sk->value.Q);
    for (int i = 0; i < maxDegree; i++) {
        parameters->ringQ->MulCoeffsBarrett(poolQP->Q, sk->value.Q, poolQP->Q); // s^2
        // 生成relin密钥
        PrivateGenSwitchingKey(poolQP->Q, &sk->value, rlk->keys[i]);
    }
    return rlk;
}

RotationKeySet *KeyGenerator::GenRotationKeysForInnerSum(SecretKey *sk)
{
    std::vector<uint64_t> tmp = parameters->GaloisElementsForRowInnerSum();
    return GenRotationKeys(tmp, sk);
}
/*
 * @brief 创建切换密钥
 *
 * @param skInput 输入密钥
 * @param skOutput 输出密钥
 * @return SwitchingKey* 返回创建的切换密钥
 *
 * 该函数:
 * - 检查P基是否为空
 * - 创建一个新的切换密钥
 * - 如果输入密钥的Q基系数大于输出密钥的Q基系数,则进行维度映射
 * - 调用PrivateGenSwitchingKey进行密钥切换
 */
SwitchingKey *KeyGenerator::GenSwitchingKey(SecretKey *skInput, SecretKey *skOutput)
{
    cout << "GenSwitchingKey" << endl;
    if (parameters->PCount() == 0) {
        throw std::runtime_error("Cannot GenSwitchingKey: modulus P is empty");
    }

    SwitchingKey *swk = new SwitchingKey(parameters, skOutput->value.Q->Level(), skOutput->value.P->Level());

    if (skInput->value.Q->Coeffs[0].size() > skOutput->value.Q->Coeffs[0].size()) {
        MapSmallDimensionToLargerDimensionNTT(skOutput->value.Q, poolQP->Q);
        MapSmallDimensionToLargerDimensionNTT(skOutput->value.P, poolQP->P);
        PrivateGenSwitchingKey(skInput->value.Q, poolQP, swk);
    } else { // N -> N or n -> N
        MapSmallDimensionToLargerDimensionNTT(skInput->value.Q, poolQ);

        if (skInput->value.Q->Level() < skOutput->value.Q->Level()) {
            Ring *ringQ = parameters->ringQ;

            ringQ->InvNTTLvl(0, poolQ, poolQP->Q);

            uint64_t Q = ringQ->Modulus[0];
            uint64_t QHalf = Q >> 1;

            Poly *polQ = poolQP->Q;
            Poly *polP = poolQ;
            uint64_t sign = 0;
            for (uint64_t j = 0; j < ringQ->N; j++) {
                uint64_t coeff = polQ->Coeffs[0][j];

                sign = 1;
                if (coeff > QHalf) {
                    coeff = Q - coeff;
                    sign = 0;
                }

                for (uint64_t i = skInput->value.Q->Level() + 1; i < skOutput->value.Q->Level() + 1; i++) {
                    polP->Coeffs[i][j] = ((coeff * sign)) | ((ringQ->Modulus[i] - coeff) * (sign ^ 1));
                }
            }

            for (uint64_t i = skInput->value.Q->Level() + 1; i < skOutput->value.Q->Level() + 1; i++) {
                ringQ->iINTTAndEqual(polP->Coeffs[i], i);
            }
        }

        PrivateGenSwitchingKey(poolQ, &skOutput->value, swk);
    }

    return swk;
}

PublicKey::PublicKey(Context *parameters)
{
    value.resize(2);
    value[0] = *parameters->NewRingQP()->NewPoly();
    value[1] = *parameters->NewRingQP()->NewPoly();
}

SecretKey::SecretKey(Context *parameters)
{
    value = *parameters->NewRingQP()->NewPoly();
}

SecretKey::SecretKey(const SecretKey &other)
{
    // 深拷贝 PolyQP
    if (other.value.Q) {
        value.Q = other.value.Q->CopyNew();
    } else {
        value.Q = nullptr;
    }

    if (other.value.P) {
        value.P = other.value.P->CopyNew();
    } else {
        value.P = nullptr;
    }
}

SecretKey &SecretKey::operator = (const SecretKey &other)
{
    if (this == &other)
        return *this; // 防止自我赋值

    // 释放当前资源
    delete value.Q;
    delete value.P;

    // 深拷贝 other 的资源
    if (other.value.Q) {
        value.Q = other.value.Q->CopyNew();
    } else {
        value.Q = nullptr;
    }

    if (other.value.P) {
        value.P = other.value.P->CopyNew();
    } else {
        value.P = nullptr;
    }

    return *this;
}
/*
 * @brief 创建切换密钥
 *
 * @param parameters 上下文参数
 * @param levelQ Q基的层级
 * @param levelP P基的层级
 *
 * 该函数:
 * - 计算分解大小
 * - 为每个分解块创建两个QP多项式
 * - 将这些多项式存储在value容器中
 */
SwitchingKey::SwitchingKey(Context *parameters, int levelQ, int levelP)
{
    int decompSize = int(ceil(double(levelQ + 1) / double(levelP + 1)));
    value.resize(decompSize);
    for (int i = 0; i < decompSize; i++) {
        value[i].resize(2);
        value[i][0] = parameters->NewRingQP()->NewPolyLvl(levelQ, levelP);
        value[i][1] = parameters->NewRingQP()->NewPolyLvl(levelQ, levelP);
    }
}
/*
 * @brief 创建旋转密钥集
 *
 * @param parameters 上下文参数
 * @param galoisElement 伽罗瓦元素列表
 *
 * 该函数:
 * - 遍历给定的伽罗瓦元素列表
 * - 为每个伽罗瓦元素创建一个新的切换密钥
 * - 将这些切换密钥存储在keys容器中
 *
 * @return 返回创建的旋转密钥集
 */
RotationKeySet::RotationKeySet(Context *parameters, const std::vector<uint64_t> &galoisElement)
{
    for (int i = 0; i < galoisElement.size(); i++) {
        uint64_t galEl = galoisElement[i];
        keys[galEl] = new SwitchingKey(parameters, parameters->QCount() - 1, parameters->PCount() - 1);
    }
}
/*
 * @brief 获取特定伽罗瓦元素的旋转切换密钥
 *
 * @param galoisEl 伽罗瓦元素
 * @return SwitchingKey* 返回与给定伽罗瓦元素对应的切换密钥
 *
 * 该函数:
 * - 检查是否存在与给定伽罗瓦元素对应的切换密钥
 * - 如果存在,则返回该切换密钥
 * - 否则,返回nullptr
 */
SwitchingKey *RotationKeySet::GetRotationKey(uint64_t galoisEl)
{
    if (keys.size() == 0) {
        return nullptr;
    }
    if (keys.find(galoisEl) != keys.end()) {
        return keys[galoisEl];
    }
    return nullptr;
}

RelinearizationKey::RelinearizationKey(Context *parameters, int maxRelinDegree)
{
    keys.resize(maxRelinDegree);
    for (int d = 0; d < maxRelinDegree; d++) {
        keys[d] = new SwitchingKey(parameters, parameters->QCount() - 1, parameters->PCount() - 1);
    }
}

EvaluationKey::EvaluationKey(RelinearizationKey *rlk, RotationKeySet *rtks) : rlk(rlk), rtks(rtks) {}

GaloisKeys::GaloisKeys(Context *parameters, const std::vector<uint64_t> &galoisElement)
{
    for (int i = 0; i < galoisElement.size(); i++) {
        uint64_t galEl = galoisElement[i];
        keys[galEl] = new SwitchingKey(parameters, parameters->QCount() - 1, parameters->PCount() - 1);
    }
}
/*
 * @brief 获取特定伽罗瓦元素的切换密钥
 *
 * @param galoisEl 伽罗瓦元素
 * @return SwitchingKey* 返回与给定伽罗瓦元素对应的切换密钥
 *
 * 该函数:
 * - 检查是否存在与给定伽罗瓦元素对应的切换密钥
 * - 如果存在,则返回该切换密钥
 * - 否则,返回nullptr
 */
SwitchingKey *GaloisKeys::GetGaloisKey(uint64_t galoisEl)
{
    if (keys.size() == 0) {
        return nullptr;
    }
    if (keys.find(galoisEl) != keys.end()) {
        return keys[galoisEl];
    }
    return nullptr;
}